from pathlib import Path
from argparse import ArgumentParser

import numpy as np
import matplotlib.pyplot as plt

from arip2.ARIP_main import load_atom_models, parse_atom_model
from arip2.typing import DTYPE


def vis_xyz(args):
    points = np.loadtxt(args.f, dtype=DTYPE)
    print('points.shape:', points.shape)
    p0 = points[0]
    print('|p0|:', np.linalg.norm(p0))

    ax = plt.axes(projection='3d')
    ax.scatter3D(points[:, 0], points[:, 1], points[:, 2])
    plt.show()


def vis_pdb(args):
    for idx, a_model in load_atom_models(args.f):
        df = parse_atom_model(a_model)
    
    ax = plt.axes(projection='3d')
    ax.scatter3D(df['x'], df['y'], df['z'])
    plt.show()


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('input', type=Path, help='*.pdb or *.xyz file')
    args = parser.parse_args()

    fp: Path = args.f
    if   fp.suffix == '.pdb': vis_pdb(args)
    elif fp.suffix == '.xyz': vis_xyz(args)
    else: raise ValueError(f'unknown suffix: {fp.suffix}')
